import jax
from jax import random
import flax
from flax import linen as nn
from typing import Sequence


class SimpleMLP(nn.Module):
    features: Sequence[int]

    @nn.compact
    def __call__(self, inputs):
        x = inputs
        for i, feat in enumerate(self.features):
            x = nn.Dense(feat, name=f'layers_{i}')(x)  # Submodule Dense must be defined in `setup()` or in a method wrapped in `@compact`
            if i != len(self.features) - 1:
                x = nn.relu(x)
        return x


key = random.PRNGKey(0)
key1, key2 = random.split(key, 2)

x = random.uniform(key1, (4, 4))
model = SimpleMLP(features=[3, 4, 5])
params = model.init(key2, x)
y = model.apply(params, x)

print('init params', jax.tree_util.tree_map(lambda x: x.shape, params))
print('x', x)
print('y', y)
